#include "cipher_mpi_base.hpp"

double bcast_pubkey(const helib::Context &context, helib::PubKey *&pk, int root,
                    MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return 0.0;
  }
  int rank;
  MPI_Comm_rank(comm, &rank);

  std::string pk_str;
  std::stringstream pk_ss;
  if (rank == root) {
    pk->writeTo(pk_ss);
    pk_str = pk_ss.str();
  }
  MPI_Bcast_String(pk_str, root, comm);
  if (rank != root) {
    pk_ss << pk_str;
    pk = new helib::PubKey(helib::PubKey::readFrom(pk_ss, context));
  }
  return pk_str.size() / 1024.0 / 1024.0;
}

void bcast_seckey(const helib::Context &context, helib::SecKey *&sk, int root,
                  MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return;
  }
  int rank;
  MPI_Comm_rank(comm, &rank);

  std::string sk_str;
  std::stringstream sk_ss;
  if (rank == root) {
    sk->writeTo(sk_ss);
    sk_str = sk_ss.str();
  }
  MPI_Bcast_String(sk_str, root, comm);
  if (rank != root) {
    sk_ss << sk_str;
    sk = new helib::SecKey(helib::SecKey::readFrom(sk_ss, context));
  }
}

void ctxt_ireduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                      int root, MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return;
  }
  int rank, size;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &size);
  auto get_idx = [=]() -> int { return (rank - root + size) % size; };
  auto get_rank = [=](int idx) -> int { return (idx + root) % size; };

  std::vector<MPI_Request> reqs(size, MPI_REQUEST_NULL);
  auto local = send_ctxt;
  std::string local_str;
  auto bound = 1l << std::max(0, static_cast<int>(NTL::NumBits(size - 1) - 1));
  for (auto stride = 1; stride <= bound; stride *= 2) {
    const int idx = get_idx();
    if (idx % stride != 0) {
      break;
    }
    bool is_recv = (idx % (2 * stride)) == 0;
    if (is_recv) {
      const int sender_idx = idx ^ stride;
      if (sender_idx < size) {
        int sender = get_rank(sender_idx);
        std::string recv_ctxt_str;
        std::stringstream ss;
        MPI_Recv_String(recv_ctxt_str, sender, 0, comm);
        ss << recv_ctxt_str;
        auto recv_add_ctxt = helib::Ctxt::readFrom(ss, local.getPubKey());
        local.addCtxt(recv_add_ctxt);
      }
    } else {
      int recipient_idx = idx ^ stride;
      int recipient = get_rank(recipient_idx);
      std::stringstream ss;
      local.writeTo(ss);
      local_str = ss.str();
      MPI_Isend_String(local_str, recipient, 0, comm, &reqs[idx]);
    }
  }
  MPI_Waitall(size, reqs.data(), MPI_STATUSES_IGNORE);
  if (rank == root) {
    *recv_ctxt = std::move(local);
  }
}

void ctxt_allreduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                        MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return;
  }
  int rank, size;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &size);
  if (size == 1) {
    return;
  }

  assert((size & (size - 1)) == 0);

  auto local = send_ctxt;
  auto bound = 1l << std::max(0, static_cast<int>(NTL::NumBits(size - 1) - 1));
  for (auto stride = 1; stride <= bound; stride *= 2) {
    const int other = rank ^ stride;
    std::string ctxt_str;
    std::stringstream ss;
    local.writeTo(ss);
    ctxt_str = ss.str();
    MPI_Isendrecv_String_replace(ctxt_str, other, 0, other, 0, comm);
    ss.str("");
    ss << ctxt_str;
    auto recv_add_ctxt = helib::Ctxt::readFrom(ss, local.getPubKey());
    local.addCtxt(recv_add_ctxt);
  }
  *recv_ctxt = std::move(local);
}

void ctxt_reduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                     int root, MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return;
  }
  int rank, size;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &size);
  auto get_idx = [=]() -> int { return (rank - root + size) % size; };
  auto get_rank = [=](int idx) -> int { return (idx + root) % size; };

  auto local = send_ctxt;
  auto bound = 1l << (NTL::NumBits(size - 1) - 1);
  for (auto stride = 1; stride <= bound; stride *= 2) {
    const int idx = get_idx();
    MPI_Comm reduce_comm;
    const auto color = (idx % stride) == 0 ? 0 : MPI_UNDEFINED;
    MPI_Comm_split(comm, color, idx, &reduce_comm);
    if (reduce_comm == MPI_COMM_NULL) {
      continue;
    }
    bool is_recv = (idx % (2 * stride)) == 0;
    if (is_recv) {
      const int sender_idx = idx ^ stride;
      if (sender_idx < size) {
        int sender = get_rank(sender_idx);
        std::string recv_ctxt_str;
        std::stringstream ss;
        MPI_Recv_String(recv_ctxt_str, sender, 0, comm);
        ss << recv_ctxt_str;
        auto recv_add_ctxt = helib::Ctxt::readFrom(ss, local.getPubKey());
        local.addCtxt(recv_add_ctxt);
      }
    } else {
      int recipient_idx = idx ^ stride;
      int recipient = get_rank(recipient_idx);
      std::stringstream ss;
      local.writeTo(ss);
      MPI_Send_String(ss.str(), recipient, 0, comm);
    }
    MPI_Barrier(reduce_comm);
    MPI_Comm_free(&reduce_comm);
  }
  if (rank == root) {
    *recv_ctxt = std::move(local);
  }
}
